# -*- coding: utf-8 -*-

import torch
import torch.nn as nn
import distiller

from config.defaults import CONFIG as cfg
from data.build_data_loader import build_data_loader
from models import resnet20


def main():
    device = cfg.DEVICE
    cfg.PHASE = "train"
    train_data_loader = build_data_loader(cfg)

    model = resnet20()
    model.to(device)
    optimizer = torch.optim.SGD(model.parameters(),
                                lr=cfg.OPTIMIZER.LEARNING_RATE,
                                momentum=0.9,
                                weight_decay=cfg.OPTIMIZER.WEIGHT_DECAY)
    criterion = nn.CrossEntropyLoss()
    criterion.to(device)

    compression_scheduler = distiller.file_config(model, optimizer, "./resnet20_slimming.yaml")

    start_epoch, epochs = cfg.TRAIN.START_EPOCH, cfg.TRAIN.EPOCHS
    for epoch in range(start_epoch, epochs + 1):
        compression_scheduler.on_epoch_begin(epoch, optimizer)
        model.train()
        for step, (image, target) in enumerate(train_data_loader, 1):
            compression_scheduler.on_minibatch_begin(epoch, step, len(train_data_loader), optimizer)

            image, target = image.to(device), target.to(device)
            predict = model(image)
            loss = criterion(predict, target)

            compression_scheduler.before_backward_pass(epoch, step, len(train_data_loader), loss,
                                                       optimizer=optimizer, return_loss_components=True)

            optimizer.zero_grad()
            loss.backward()

            compression_scheduler.before_parameter_optimization(epoch, step, len(train_data_loader), optimizer)

            optimizer.step()

            compression_scheduler.on_minibatch_end(epoch, step, len(train_data_loader), optimizer)
            break
        compression_scheduler.on_epoch_end(epoch, optimizer)
        break


if __name__ == "__main__":
    main()
